//! RK3588 I2C Controller Driver
//! 
//! This module provides a driver for the RK3588 I2C controller.
//! 
//! # Hardware Specifications
//! - RK3588 has 9 I2C controllers (I2C0-I2C8)
//! - I2C0-I2C5: Support standard mode (100Kbps) and fast mode (400Kbps)
//! - I2C6-I2C8: Support standard mode, fast mode, and fast mode plus (1Mbps)
//! 
//! # Memory Map
//! - I2C0: 0xFD88_0000
//! - I2C1: 0xFEA9_0000
//! - I2C2: 0xFEA9_4000
//! - I2C3: 0xFEA9_8000
//! - I2C4: 0xFEA9_C000
//! - I2C5: 0xFEAA_0000
//! - I2C6: 0xFEA5_0000
//! - I2C7: 0xFEA7_0000
//! - I2C8: 0xFEA8_0000

use core::ptr::{read_volatile, write_volatile};

/// I2C Register Offsets
const I2C_CON: usize = 0x00;      // Control Register
const I2C_CLKDIV: usize = 0x04;   // Clock Divisor Register
const I2C_MRXADDR: usize = 0x08;  // Master RX Address Register
const I2C_MRXRADDR: usize = 0x0C; // Master RX Register Address
const I2C_MTXCNT: usize = 0x10;   // Master TX Count Register
const I2C_MRXCNT: usize = 0x14;   // Master RX Count Register
const I2C_IEN: usize = 0x18;      // Interrupt Enable Register
const I2C_IPD: usize = 0x1C;      // Interrupt Pending Register
const I2C_FCNT: usize = 0x20;     // Finished Count Register
const I2C_TXDATA: [usize; 8] = [  // TX Data Registers
    0x100, 0x104, 0x108, 0x10C, 0x110, 0x114, 0x118, 0x11C
];
const I2C_RXDATA: [usize; 8] = [  // RX Data Registers
    0x200, 0x204, 0x208, 0x20C, 0x210, 0x214, 0x218, 0x21C
];

/// I2C Control Register Bits
const I2C_CON_EN: u32 = 1 << 0;           // I2C Enable
const I2C_CON_MODE_MASK: u32 = 0x3 << 1;  // Mode Mask
const I2C_CON_MODE_TX: u32 = 0x0 << 1;    // Transmit Mode
const I2C_CON_MODE_RX: u32 = 0x1 << 1;    // Receive Mode
const I2C_CON_START: u32 = 1 << 3;        // Start Condition
const I2C_CON_STOP: u32 = 1 << 4;         // Stop Condition
const I2C_CON_NAKSTOP: u32 = 1 << 6;      // NAK Stop Enable
const I2C_CON_ACTACK: u32 = 1 << 5;       // ACK Active

/// I2C Interrupt Bits
const I2C_INT_MBTF: u32 = 1 << 0;   // Master Byte Transfer Finished
const I2C_INT_MBRF: u32 = 1 << 1;   // Master Byte Receive Finished
const I2C_INT_MBRFIEN: u32 = 1 << 2; // Master Byte Receive FIFO Interrupt Enable
const I2C_INT_NAKRCV: u32 = 1 << 3; // NAK Received
const I2C_INT_STOP: u32 = 1 << 4;   // Stop Condition
const I2C_INT_START: u32 = 1 << 5;  // Start Condition

/// I2C Base Addresses
#[derive(Debug, Clone, Copy)]
pub enum I2cController {
    I2C0 = 0xFD88_0000,
    I2C1 = 0xFEA9_0000,
    I2C2 = 0xFEA9_4000,
    I2C3 = 0xFEA9_8000,
    I2C4 = 0xFEA9_C000,
    I2C5 = 0xFEAA_0000,
    I2C6 = 0xFEA5_0000,
    I2C7 = 0xFEA7_0000,
    I2C8 = 0xFEA8_0000,
}

/// I2C Speed Mode
#[derive(Debug, Clone, Copy)]
pub enum I2cSpeed {
    Standard,  // 100 Kbps
    Fast,      // 400 Kbps
    FastPlus,  // 1 Mbps
}

/// I2C Error Types
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum I2cError {
    Timeout,
    Nack,
    BusError,
    InvalidParameter,
}

/// RK3588 I2C Driver
pub struct Rk3588I2c {
    base_addr: usize,
    speed: I2cSpeed,
}

impl Rk3588I2c {
    /// Create a new I2C controller instance
    pub fn new(controller: I2cController, speed: I2cSpeed) -> Self {
        let mut i2c = Self {
            base_addr: controller as usize,
            speed,
        };
        i2c.init();
        i2c
    }

    /// Initialize the I2C controller
    fn init(&mut self) {
        // Disable I2C first
        self.write_reg(I2C_CON, 0);

        // Set clock divisor based on speed
        // Assuming I2C clock is 200MHz
        let clk_div = match self.speed {
            I2cSpeed::Standard => 1000,  // ~100 Kbps
            I2cSpeed::Fast => 250,       // ~400 Kbps
            I2cSpeed::FastPlus => 100,   // ~1 Mbps
        };
        self.write_reg(I2C_CLKDIV, clk_div);

        // Enable NAK stop
        self.write_reg(I2C_CON, I2C_CON_NAKSTOP | I2C_CON_EN);

        // Clear all interrupts
        self.write_reg(I2C_IPD, 0x7F);
    }

    /// Write to register
    #[inline]
    fn write_reg(&self, offset: usize, value: u32) {
        unsafe {
            write_volatile((self.base_addr + offset) as *mut u32, value);
        }
    }

    /// Read from register
    #[inline]
    fn read_reg(&self, offset: usize) -> u32 {
        unsafe { read_volatile((self.base_addr + offset) as *const u32) }
    }

    /// Wait for transfer to complete
    fn wait_complete(&self, timeout_ms: u32) -> Result<(), I2cError> {
        let timeout_cycles = timeout_ms * 24_000; // Assuming 24MHz timer
        let mut cycles = 0;

        while cycles < timeout_cycles {
            let ipd = self.read_reg(I2C_IPD);
            
            // Check for NAK
            if ipd & I2C_INT_NAKRCV != 0 {
                self.write_reg(I2C_IPD, I2C_INT_NAKRCV);
                return Err(I2cError::Nack);
            }

            // Check for completion
            if ipd & (I2C_INT_MBTF | I2C_INT_STOP) != 0 {
                self.write_reg(I2C_IPD, ipd);
                return Ok(());
            }

            cycles += 1;
        }

        Err(I2cError::Timeout)
    }

    /// Write data to I2C device
    pub fn write(&mut self, addr: u8, data: &[u8]) -> Result<(), I2cError> {
        if data.is_empty() || data.len() > 32 {
            return Err(I2cError::InvalidParameter);
        }

        // Clear interrupts
        self.write_reg(I2C_IPD, 0x7F);

        // Set slave address (7-bit addressing)
        self.write_reg(I2C_MRXADDR, (addr << 1) as u32);

        // Set transmit count
        self.write_reg(I2C_MTXCNT, data.len() as u32);

        // Write data to TX FIFO
        let mut word_idx = 0;
        let mut byte_idx = 0;
        while byte_idx < data.len() {
            let mut word = 0u32;
            for i in 0..4 {
                if byte_idx + i < data.len() {
                    word |= (data[byte_idx + i] as u32) << (i * 8);
                }
            }
            self.write_reg(I2C_TXDATA[word_idx], word);
            word_idx += 1;
            byte_idx += 4;
        }

        // Start transmission
        let con = I2C_CON_EN | I2C_CON_MODE_TX | I2C_CON_START | I2C_CON_STOP;
        self.write_reg(I2C_CON, con);

        // Wait for completion
        self.wait_complete(100)?;

        Ok(())
    }

    /// Read data from I2C device
    pub fn read(&mut self, addr: u8, buffer: &mut [u8]) -> Result<(), I2cError> {
        if buffer.is_empty() || buffer.len() > 32 {
            return Err(I2cError::InvalidParameter);
        }

        // Clear interrupts
        self.write_reg(I2C_IPD, 0x7F);

        // Set slave address (7-bit addressing, read bit set)
        self.write_reg(I2C_MRXADDR, ((addr << 1) | 0x01) as u32);

        // Set receive count
        self.write_reg(I2C_MRXCNT, buffer.len() as u32);

        // Start reception
        let con = I2C_CON_EN | I2C_CON_MODE_RX | I2C_CON_START | I2C_CON_STOP | I2C_CON_ACTACK;
        self.write_reg(I2C_CON, con);

        // Wait for completion
        self.wait_complete(100)?;

        // Read data from RX FIFO
        let mut word_idx = 0;
        let mut byte_idx = 0;
        while byte_idx < buffer.len() {
            let word = self.read_reg(I2C_RXDATA[word_idx]);
            for i in 0..4 {
                if byte_idx + i < buffer.len() {
                    buffer[byte_idx + i] = ((word >> (i * 8)) & 0xFF) as u8;
                }
            }
            word_idx += 1;
            byte_idx += 4;
        }

        Ok(())
    }

    /// Write data to a specific register of I2C device
    pub fn write_reg_data(&mut self, addr: u8, reg: u8, data: &[u8]) -> Result<(), I2cError> {
        if data.len() > 31 {
            return Err(I2cError::InvalidParameter);
        }

        // Prepare buffer with register address + data
        let mut buffer = [0u8; 33];
        buffer[0] = reg;
        buffer[1..1 + data.len()].copy_from_slice(data);

        self.write(addr, &buffer[..1 + data.len()])
    }

    /// Read data from a specific register of I2C device
    pub fn read_reg_data(&mut self, addr: u8, reg: u8, buffer: &mut [u8]) -> Result<(), I2cError> {
        // Write register address
        self.write(addr, &[reg])?;
        
        // Read data
        self.read(addr, buffer)?;

        Ok(())
    }
}

/// Global I2C instance helper
static mut I2C_INSTANCES: [Option<Rk3588I2c>; 9] = [None, None, None, None, None, None, None, None, None];

/// Get or initialize an I2C controller instance
pub fn get_i2c(controller: I2cController, speed: I2cSpeed) -> &'static mut Rk3588I2c {
    let idx = match controller {
        I2cController::I2C0 => 0,
        I2cController::I2C1 => 1,
        I2cController::I2C2 => 2,
        I2cController::I2C3 => 3,
        I2cController::I2C4 => 4,
        I2cController::I2C5 => 5,
        I2cController::I2C6 => 6,
        I2cController::I2C7 => 7,
        I2cController::I2C8 => 8,
    };

    unsafe {
        if I2C_INSTANCES[idx].is_none() {
            I2C_INSTANCES[idx] = Some(Rk3588I2c::new(controller, speed));
        }
        I2C_INSTANCES[idx].as_mut().unwrap()
    }
}

